"""
此处假定最理想的情况下进行推理代码的编写，此版本为alpha版本
    1. 输入图像经过resize = patch_size
    2. 图像的边界没有黑框，也就是不用进行crop
    3. 由于目标等于patch_size，所以也没有sliding_windows
        - 因此也不用使用高斯核融合多次计算的结果

TODO:
    还需要找出训练参数，用于更好的描述限定条件
    路老哥训练使用的命令:
        - predict:nnUNet_predict -i /home/imed/LJY/OCT_nnUNet/Task089_PapSmear/imagesTs -o /home/imed/LJY/OCT_nnUNet/predicted_10 -t 89 -m 2d -f 4
        - plan and process: nnUNet_plan_and_preprocess -t 89 -pl3d None --verify_dataset_integrity
        - train:CUDA_VISIBLE_DEVICES=1 nnUNet_train 2d nnUNetTrainerV2 89 4

"""
import time

import cv2
import numpy
import torch
import torch.nn.functional as F

model = torch.jit.load("v1_jit.pth")
model.eval()
# load data
sample_image = "../resource/raw_resize/1.jpg"
start_time = time.perf_counter()
mat = cv2.imread(sample_image)
mat = mat.transpose((2, 0, 1))
mat = mat.astype(numpy.float32)
for c in range(mat.shape[0]):
    mat[c] = (mat[c] - mat[c].mean()) / (mat[c].std() + 1e-8)
mat = numpy.expand_dims(mat, 0)
# infer
with torch.no_grad():
    x = torch.from_numpy(mat)
    pred = F.softmax(model(x), 1)
    pred += torch.flip(F.softmax(model(torch.flip(x, (3,))), 1), (3,))
    pred += torch.flip(F.softmax(model(torch.flip(x, (2,))), 1), (2,))
    pred += torch.flip(F.softmax(model(torch.flip(x, (3, 2))), 1), (3, 2))
    pred /= 4
    pred = pred.argmax(1)
    # save
    pred = pred.cpu().numpy()[0]
    cv2.imwrite("out_pred_torchscript.png", pred * 255)
end_time = time.perf_counter()
print("time cost: ", end_time - start_time, "s")